#include "nn.h"
#include "dataloader.h"
#include "mat.h"
#include "preprocess.h"
#include <utility>
using namespace std;


class Model:public nn
{
    public:
        Conv2d* conv1;
        MaxPool* maxpool1;
        Relu* relu1;
        Conv2d* conv2;
        MaxPool* maxpool2;
        Relu* relu2;
        Flatten* flatten;
        Linear* l1;
        Linear* l2;
        Softmax* softmax;

        Model()
        {
            this->conv1=nn::conv2d(3, 32, 3, 1);
            this->maxpool1=nn::maxPool(32, 2, 2);
            this->relu1=nn::relu();
            this->conv2=nn::conv2d(32, 64, 3, 1);
            this->maxpool2=nn::maxPool(64, 2, 2);
            this->relu2=nn::relu();
            this->flatten=nn::flatten();
            this->l1=nn::linear(4096, 128);
            this->l2=nn::linear(128, 10);
            this->softmax=nn::softmax();
        }

        ~Model()
        {
            delete this->conv1;
            delete this->maxpool1;
            delete this->relu1;
            delete this->conv2;
            delete this->maxpool2;
            delete this->relu2;
            delete this->flatten;
            delete this->l1;
            delete this->l2;
            delete this->softmax;
        }

        FloatMat* forward(FloatMat* input)
        {
            FloatMat* output1=this->conv1->forward(input);
            FloatMat* output2=this->maxpool1->forward(output1);
            FloatMat* output3=this->relu1->forward(output2);
            FloatMat* output4=this->conv2->forward(output3);
            FloatMat* output5=this->maxpool2->forward(output4);
            FloatMat* output6=this->relu2->forward(output5);
            FloatMat* output7=this->flatten->forward(output6);
            FloatMat* output8=this->l1->forward(output7);
            FloatMat* output9=this->l2->forward(output8);
            FloatMat* output10=this->softmax->forward(output9);
            return output10;
        }
};

void dataTrain(Model* model, Dataloader dataloaderTrain, int epochNum, int batch, float learningRate)
{
    model->setWithGrad();
    model->setLearningRate(learningRate);
    CrossEntropyLoss* loss=new CrossEntropyLoss();
    SGDOptimizer* optimizer = new SGDOptimizer(learningRate);
    model->setOptimizer(optimizer);
    for(int epoch=0;epoch<epochNum;epoch++)
    {
        cout << "Epoch:" << epoch+1 << endl;
        for(int i=0;i<dataloaderTrain.imgs.size()/batch;i++)
        {
            float lossSum=0;
            model->zeroGrad();
            for(int j=0;j<batch;j++)
            {
                pair<FloatMat*, int> p=dataloaderTrain.getARandomImgMat();
                FloatMat* label=call(FloatMat::getSpecialFloatMat, 1, 1, 10, 0);
                label->values[0][0][p.second]=1;
                FloatMat* output=model->forward(p.first);
                lossSum += loss->getLoss(output, label);
                model->backward(loss->grad);
                delete p.first;
                delete label;
            }
            model->step();
            cout << lossSum/batch << endl;
        }
    }
    delete loss;
    delete optimizer;
    monitor.openMonitor();
    monitor.checkMat();
}

void dataTest(Model* model, Dataloader dataloaderTest, int batch)
{
    cout << "------------verification-----------" << endl;
    model->setNoGrad();
    CrossEntropyLoss* loss = new CrossEntropyLoss();
    int correctNum = 0;
    for(int i=0;i<dataloaderTest.imgs.size()/batch;i++)
    {
        float sumLoss = 0;
        for(int j=0;j<batch;j++)
        {
            pair<FloatMat*, int> p = dataloaderTest.getARandomImgMat();
            FloatMat* label = call(FloatMat::getSpecialFloatMat, 1, 1, 10, 0);
            label->values[0][0][p.second] = 1;
            FloatMat* output = model->forward(p.first);
            sumLoss += loss->getLoss(output, label);
            if(FloatMat::argMax(label) == p.second) correctNum++;
            delete p.first;
            delete label;
        }
        cout << i+1 << ": Loss:" << sumLoss/batch << endl;
    }
    cout << "Accuracy:" << correctNum/dataloaderTest.imgs.size() << endl;
    gcMat(loss);
    monitor.openMonitor();
    monitor.checkMat();
}

int main()
{
    float learningRate=0.1;
    Dataloader dataloaderTrain;
    Dataloader dataloaderTest;
    dataloaderTrain.loadAllImgsFromPath("../data/cifar10/train", regex(R"((\d)_(\d+))"));
    dataloaderTest.loadAllImgsFromPath("../data/cifar10/test", regex(R"((\d)_(\d+))"));
    Model* model=new Model();
    dataTrain(model, dataloaderTrain, 10, 64, learningRate);
    dataTest(model, dataloaderTest, 64);
    model->save("./model.pytm");
    delete model;
    return 0;
}
